from dolfin import *
import numpy

parameters['allow_extrapolation'] = True

'''can I have a weirdly shaped mesh created by hand and a function space on it?'''
mesh = Mesh()
editor = MeshEditor()
editor.open(mesh,2,2)

editor.init_vertices(8)
editor.init_cells(6)

editor.add_vertex(0,0,0)
editor.add_vertex(1,0.5,0)
editor.add_vertex(2,0,0.5)
editor.add_vertex(3,0.5,0.5)
editor.add_vertex(4,1,0.5)
editor.add_vertex(5,0,1)
editor.add_vertex(6,0.5,1)
editor.add_vertex(7,1,1)

editor.add_cell(0,0,1,3)
editor.add_cell(1,0,2,3)
editor.add_cell(2,2,3,6)
editor.add_cell(3,2,5,6)
editor.add_cell(4,3,4,7)
editor.add_cell(5,3,6,7)
editor.close()

mesh.order()
V = FunctionSpace(mesh,'CG',1)

'''can I create a mesh like the one above from the vertices of the UnitSquareMesh?'''
full_mesh = RectangleMesh(0,0,1,1,100,100)

H = FunctionSpace(full_mesh,'CG',1)
V = FunctionSpace(full_mesh,'DG',0)
phi = interpolate(Expression('sqrt((x[0]-0.5)*(x[0]-0.5)+(x[1]-0.5)*(x[1]-0.5))-0.2'),V)
phi_H = interpolate(Expression('sqrt((x[0]-0.5)*(x[0]-0.5)+(x[1]-0.5)*(x[1]-0.5))-0.2'),H)
plot(phi,interactive=True)
hmin = full_mesh.hmin()
values = numpy.ndarray.astype(numpy.where(numpy.abs(phi.vector().array()) < 1*hmin,1,0),'uintp')

#print values
cell_function = MeshFunction('size_t',full_mesh,full_mesh.topology().dim()) 
cell_function.set_all(0)
cell_function.set_values(values)

plot(cell_function,interactive=True)
ls_mesh = SubMesh(full_mesh,cell_function,1)
plot(ls_mesh,interactive=True)
W = FunctionSpace(ls_mesh,'DG',0)
ls_sd = interpolate(phi,W)
plot(ls_sd,interactive=True)
ls_sd.vector()[:] += 0.1
plot(ls_sd,interactive=True)

u = TrialFunction(H)
v = TestFunction(H)

a = u*v*dx(1)
l = ls_sd*v*dx(1)
A,l = assemble_system(a,l,cell_domains=cell_function)
A.ident_zeros()
u_fine = Function(H)
solve(A,u_fine.vector(),l)
plot(u_fine,interactive=True)

a = u*v*dx(0)
l = phi*v*dx(0)
A,l = assemble_system(a,l,cell_domains=cell_function)
A.ident_zeros()
u_course = Function(H)
solve(A,u_course.vector(),l)
plot(u_course,interactive=True)

u_moved = Function(H)
u_moved.vector()[:] = u_fine.vector().array() + u_course.vector().array()
plot(u_moved,interactive=True)
